#!/usr/bin/env python
# -*- coding: utf-8 -*-

from threading import Timer
import tushare as ts
import re
import aiohttp
import asyncio
from datetime import datetime, timezone, timedelta
import time
import codecs
import os
import random


ua = [
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
"Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.94 Safari/537.36",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:56.0) Gecko/20100101 Firefox/56.0",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_1) AppleWebKit/604.3.5 (KHTML, like Gecko) Version/11.0.1 Safari/604.3.5",
"Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:56.0) Gecko/20100101 Firefox/56.0",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/604.3.5 (KHTML, like Gecko) Version/11.0.1 Safari/604.3.5",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.89 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/604.1.38 (KHTML, like Gecko) Version/11.0 Safari/604.1.38"
]


class StockTicker:
    """
    实时行情报价器
    """

    _url = 'http://hq.sinajs.cn/list='

    def __init__(self, dir_path):
        self.code_list = self.get_code_list(self)  # 股票代码列表，至少每天更新一次
        self.logger = Logger(BatchFlushReciever(dir_path))  # 持久化记录
        self.timer = None

    @staticmethod
    def splist(l, s):
        """将list l按照s的长度分段"""
        return [l[i:i + s] for i in range(len(l)) if i % s == 0]

    @staticmethod
    def get_code_list(self):
        """将列表分段"""
        df = ts.get_stock_basics()
        full_list = list(map(lambda code: str('sh' + code[0:] if code[0] == '6' else 'sz' + code[0:]), list(df.index)))
        return list(map(lambda a: ','.join(a), self.splist(full_list, 400)))

    def update_stock_list(self):
        self.code_list = self.get_code_list(self)

    @staticmethod
    def get_header():
        return {
            "accept-charset": "ISO-8859-1,utf-8;q=0.7,*;q=0.3",
            "accept-language": "en-US,en;q=0.8",
            "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            "user-agent": random.choice(ua),
            "Accept-Encoding": "gzip,deflate",
        }

    @staticmethod
    async def fetch(session, url):
        """协程并发http请求"""
        with aiohttp.Timeout(10, loop=session.loop):
            async with session.get(url) as response:
                return await response.text()

    async def fetch_ticks(self, loop, url):
        """抓取tick数据"""
        async with aiohttp.ClientSession(loop=loop, headers=self.get_header()) as session:
            try:
                html = await self.fetch(session, url)
                # 将tick字段按逗号隔开
                rows = re.compile(r'(var hq_str_)|(\")|;', re.M).subn('', html)[0].replace('=', ',')
                # print(rows)
                for tick in rows.split('\n'):
                    if len(tick) > 10:
                        self.logger.tick_one(tick)
                # print(ticks)
                if self.timer and self.timer.isAlive:
                    self.timer.cancel()
                self.timer = Timer(20, self.flush)
                self.timer.start()
            except Exception as err:
                print(err)

    def run_ticker(self):
        """全量采集一次所有股票当前tick"""
        for stocks in self.code_list:
            loop = asyncio.get_event_loop()
            loop.run_until_complete(self.fetch_ticks(loop, self._url + stocks))

    def flush(self):
        self.logger.flush()
        if self.timer and self.timer.isAlive:
            self.timer.cancel()


class Logger:
    """
    数据记录类，接收全量混合的所有股票tick每只股票建立一个实体对象
    """
    def __init__(self, reciever):
        self.reciever = reciever
        self.stocks_dict = {}  # 去重hash map

    @staticmethod
    def get_code(tick):
        """获取股票代码"""
        m = re.search('s[hz][0-9]{6}', tick)
        if m:
            g = m.group()
            if isinstance(g, str):
                return g
        return None

    @staticmethod
    def get_event_time(tick):
        """获取时间"""
        m = re.search('[0-9]{4}-[0-9]{2}-[0-9]{2},[0-9]{2}:[0-9]{2}:[0-9]{2}', tick)
        if m:
            g = m.group()
            if isinstance(g, str):
                datetime.strptime(g, '%Y-%m-%d,%H:%M:%S')
                dtime = datetime.strptime(g, '%Y-%m-%d,%H:%M:%S').replace(tzinfo=timezone(timedelta(hours=8)))
                return int(time.mktime(dtime.timetuple()))
        return None

    def tick_one(self, tick):
        """记录一个tick"""
        code = self.get_code(tick)
        event_time = self.get_event_time(tick)
        tick_logger = self.stocks_dict.get(code)
        if not tick_logger:  # 未创建单条记录器
            tick_logger = TickLogger(code, self.reciever)
            self.stocks_dict[code] = tick_logger
        tick_logger.append(event_time, tick)

    def flush(self):
        for log in self.stocks_dict.values():
            log.flush()
        self.stocks_dict.clear()
        self.reciever.flush()


class BatchFlushReciever:
    def __init__(self, dir_path):
        self.dir_path = dir_path
        self.queue = []
        self.timer = None

    def append(self, item):
        self.queue.append(item)
        if len(self.queue) >= 8000:  # 缓存超出8000条记录则持久化
            self.flush()

    def flush(self):
        queue = sorted(self.queue[:])
        self.queue.clear()
        file = codecs.open(os.path.join(self.dir_path, datetime.today().strftime('%Y-%m-%d') + '.csv'), 'a', 'utf8')
        file.write('\n'.join(queue) + '\n')
        file.close()


class TickLogger:
    def __init__(self, code, reciever):
        self.code = code

        self.reciever = reciever
        self.tick_cache = {}  # 缓存堆栈

    def append(self, event_time, tick):
        """ 添加一条tick记录， 使缓冲区始终存在3条记录"""
        if not self.tick_cache.get(event_time):
            self.tick_cache[event_time] = tick
            if len(self.tick_cache) > 4:  # 缓存大于等于3，则清除最旧的一条缓存，进行持久化
                self.reciever.append(self.tick_cache.pop(sorted(self.tick_cache.keys())[0]))

    def flush(self):
        """清空缓冲区"""
        for et in sorted(self.tick_cache.keys()):
            self.reciever.append(self.tick_cache.get(et))
        self.tick_cache.clear()

ticker = StockTicker

if __name__ == '__main__':
    import inspect
    this_file = inspect.getfile(inspect.currentframe())
    dir_path = os.path.abspath(os.path.dirname(this_file))
    StockTicker(dir_path).run_ticker()
    print('running')
